import numpy as np
import matplotlib.pyplot as plt

from mpl_toolkits.mplot3d import Axes3D


def get_perceptions(position, obstacles, targets):
    obstacle_detections = []
    for obstacle in obstacles:
        distance = np.linalg.norm(position - obstacle['position'])
        if distance < obstacle['radius']:
            obstacle_detections.append(obstacle)

    target_detections = []
    for target in targets:
        distance = np.linalg.norm(position - target['position'])
        if distance < 100:
            target_detections.append(target)

    return obstacle_detections, target_detections


# 决策函数
def make_decision(obstacle_detections, target_detections):
    # 基于感知数据做出决策
    # ...

    # 更新速度和方向
    new_velocity = np.array([0, 0, 0])
    new_direction = np.array([0, 0, 0])

    return new_velocity, new_direction


# 球体函数
# 球体函数
def sphere(x, y, z, radius, segments, color):
    u = np.linspace(0, 2 * np.pi, segments)
    v = np.linspace(0, np.pi, segments)

    X = x + radius * np.cos(u) * np.sin(v)
    Y = y + radius * np.sin(u) * np.sin(v)
    Z = z + radius * np.cos(v)

    # 扩展 Z 为二维数组
    Z = np.expand_dims(Z, axis=1)

    surf = ax.plot_surface(X, Y, Z, color=color)
    return surf


# 定义障碍物和环境参数
obstacles = [
    {'type': 'building', 'position': [100, 100, 50], 'radius': 50},
    {'type': 'airflow', 'position': [200, 200, 100], 'radius': 50},
    {'type': 'gravity', 'position': [300, 300, 150], 'radius': 50}
]

# 定义目标参数
targets = [
    {'position': [400, 400, 100], 'velocity': [10, 10, 10]},
    {'position': [500, 500, 100], 'velocity': [-10, -10, -10]}
]

# 初始化飞行器状态
position = np.array([0, 0, 0])  # 位置 [x, y, z]
velocity = np.array([0, 0, 0])  # 速度 [vx, vy, vz]

# 仿真时间步长
dt = 0.1

# 仿真时间
t = 0

# 仿真数据存储
positions = []
velocities = []

# 仿真循环
while t < 100:
    # 获取障碍物和目标的感知数据
    obstacle_detections, target_detections = get_perceptions(position, obstacles, targets)

    # 基于感知数据做出决策
    new_velocity, new_direction = make_decision(obstacle_detections, target_detections)

    # 更新飞行器状态
    position = position.astype(float)
    position += velocity * dt
    velocity = new_velocity

    # 存储数据
    positions.append(position)
    velocities.append(velocity)

    # 更新时间
    t += dt

# 绘制数据可视化

# 3D 绘图
fig = plt.figure(figsize=(10, 10))
ax = fig.add_subplot(111, projection='3d')

# 绘制障碍物
for obstacle in obstacles:
    if obstacle['type'] == 'building':
        color = 'r'
    elif obstacle['type'] == 'airflow':
        color = 'g'
    elif obstacle['type'] == 'gravity':
        color = 'b'
    ax.add_artist(
        sphere(obstacle['position'][0], obstacle['position'][1], obstacle['position'][2], obstacle['radius'], 20,
               color))

# 绘制目标
for target in targets:
    ax.scatter(target['position'][0], target['position'][1], target['position'][2], color='b')

# 绘制飞行器轨迹
ax.plot(*zip(*positions), color='k')

# 绘制飞行器速度向量
ax.quiver(*zip(*positions), *zip(*velocities), color='m')

ax.set_xlabel('X')
ax.set_ylabel('Y')
ax.set_zlabel('Z')
ax.set_title('飞行器仿真')

# 2D 位置和速度图
fig, (ax1, ax2) = plt.subplots(2, 2, figsize=(10, 10), sharex=True, sharey=True)
fig, (ax3, ax4) = plt.subplots(2, 2, figsize=(10, 10), sharex=True, sharey=True)

# 位置图

fig, ax1 = plt.subplots()
fig, ax2 = plt.subplots()
fig, ax3 = plt.subplots()
fig, ax4 = plt.subplots()

# 创建一个 Line2D 对象
line1 = plt.Line2D([p[0] for p in positions], [p[1] for p in positions], color='r')
# 将 Line2D 对象添加到子图
ax1.add_line(line1)
# 设置子图属性
ax1.set_xlabel('时间')
ax1.set_ylabel('X')
ax1.set_title('位置')

line2 = plt.Line2D([p[0] for p in positions], [p[2] for p in positions], color='b')
# 将 Line2D 对象添加到子图
ax2.add_line(line2)
ax2.set_xlabel('时间')
ax2.set_ylabel('Z')

# 速度图
line3 = plt.Line2D([v[0] for v in velocities], [v[1] for v in velocities], color='r')
ax3.add_line(line3)
ax3.set_xlabel('时间')
ax3.set_ylabel('X')
ax3.set_title('速度')

ax4.plot([v[0] for v in velocities], [v[1] for v in velocities], color='b')
ax4.set_xlabel('时间')
ax4.set_ylabel('Z')

plt.show()


# 感知函数
